#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    :   pack_fw_states.py
@Time    :   2024/05/25 14:37:36
@Author  :   Lin Yuheng 
@Version :   1.0
@Desc    :   None
"""

from outdoor_offb.msg import FWstates
from fixed_wing_sub_pub import *



class PackFWStates:
    def __init__(self):

        self.planeID = 0

        self.uavID = "plane_0"

        self.print_couter = 0

        self.fixed_wing_sub_pub = FixedWingSubPub()

        # 转换四元数中间量
        self.att_angle = np.zeros(3)
        self.att_quat = np.zeros(4)

    def ros_sub_pub(self):
        ###############################订阅mavros消息+控制器指令###################################
        self.fixed_wing_states_sub = rospy.Subscriber(
            self.uavID + "/mavros/state",
            State,
            callback=self.fixed_wing_sub_pub.state_cb,
            queue_size=10,
        )

        # 【订阅】无人机imu信息，
        self.fixed_wing_imu_sub = rospy.Subscriber(
            self.uavID + "/mavros/imu/data",
            Imu,
            callback=self.fixed_wing_sub_pub.imu_cb,
            queue_size=10,
        )

        # 【订阅】无人机gps位置，
        self.fixed_wing_global_position_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/global_position/global",
            NavSatFix,
            callback=self.fixed_wing_sub_pub.global_position_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机gps相对alt
        self.fixed_wing_global_rel_alt_from_px4_cb = rospy.Subscriber(
            self.uavID + "/mavros/global_position/rel_alt",
            Float64,
            callback=self.fixed_wing_sub_pub.global_rel_alt_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机ump位置
        self.fixed_wing_ump_position_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/global_position/local",
            Odometry,
            callback=self.fixed_wing_sub_pub.umt_position_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机gps三向速度
        self.fixed_wing_velocity_global_fused_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/global_position/raw/gps_vel",
            TwistStamped,
            callback=self.fixed_wing_sub_pub.velocity_global_fused_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机ned位置
        self.fixed_wing_local_position_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/local_position/pose",
            PoseStamped,
            callback=self.fixed_wing_sub_pub.local_position_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机ned三向速度
        self.fixed_wing_local_velocity_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/local_position/velocity_local",
            TwistStamped,
            callback=self.fixed_wing_sub_pub.velocity_ned_fused_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机ned三向加速度
        self.fixed_wing_ned_acceleration_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/local_position/accel",
            AccelWithCovarianceStamped,
            callback=self.fixed_wing_sub_pub.acc_ned_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】风速估计
        self.fixed_wing_wind_estimate_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/wind_estimation",
            TwistWithCovarianceStamped,
            callback=self.fixed_wing_sub_pub.wind_estimate_from_px4_cb,
            queue_size=10,
        )

        # # 【订阅】电池状态
        # self.fixed_wing_battery_states_from_px4_sub = rospy.Subscriber(
        #     "mavros/battery",
        #     BatteryState,
        #     callback=self.fixed_wing_sub_pub.battery_state_from_px4_cb,
        #     queue_size=10,
        # )

        # 【订阅】无人机当前航点
        self.fixed_wing_waypoints_cb = rospy.Subscriber(
            self.uavID + "/mavros/mission/waypoints",
            WaypointList,
            callback=self.fixed_wing_sub_pub.waypointlist_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机到达的航点
        self.fixed_wing_waypoint_reached_cb = rospy.Subscriber(
            self.uavID + "/mavros/mission/reached",
            WaypointReached,
            callback=self.fixed_wing_sub_pub.waypoints_reached_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机当前的高度
        self.fixed_wing_altitude_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/altitude",
            Altitude,
            callback=self.fixed_wing_sub_pub.altitude_from_px4_cb,
            queue_size=10,
        )

        # 【订阅】无人机当前的空速，地速
        self.fixed_wing_air_ground_speed_from_px4_sub = rospy.Subscriber(
            self.uavID + "/mavros/vfr_hud",
            VFR_HUD,
            callback=self.fixed_wing_sub_pub.air_ground_speed_from_px4_cb,
            queue_size=10,
        )

        # # 【订阅】上层控制器的四通道控制量
        # self.fixed_wing_cmd_from_controller_sub = rospy.Subscriber(
        #     self.uavID + "/fixed_wing_control/fw_cmd",
        #     FWcmd,
        #     callback=self.fixed_wing_sub_pub.cmd_from_controller_cb,
        #     queue_size=10,
        # )

        ###############################发布mavros消息+飞机状态###################################
        # self.fixed_wing_local_pos_sp_pub = rospy.Publisher(
        #     self.uavID + "/mavros/setpoint_position/local", PoseStamped, queue_size=10
        # )

        # self.fixed_wing_global_pos_sp_pub = rospy.Publisher(
        #     self.uavID + "/mavros/setpoint_raw/global", GlobalPositionTarget, queue_size=10
        # )

        # self.fixed_wing_local_att_sp_pub = rospy.Publisher(
        #     self.uavID + "/mavros/setpoint_raw/attitude", AttitudeTarget, queue_size=10
        # )

        self.fixed_wing_states_pub = rospy.Publisher(
            self.uavID + "/fixed_wing_control/fw_state", FWstates, queue_size=10
        )

        ###############################发布mavros消息+飞机状态###################################

        self.set_mode_service = rospy.ServiceProxy(
            self.uavID + "/mavros/set_mode", SetMode
        )

        self.arming_service = rospy.ServiceProxy(
            self.uavID + "/mavros/cmd/arming", CommandBool
        )

    def pack_fw_states(self):
        # 给结构体赋值，更新飞机状态
        self.fixed_wing_sub_pub.fw_states_from_mavros.planeID = (
            self.planeID
        )
        # 控制模式
        self.fixed_wing_sub_pub.fw_states_from_mavros.control_mode = (
            self.fixed_wing_sub_pub.current_state.mode
        )

        # 以下为GPS信息
        self.fixed_wing_sub_pub.fw_states_from_mavros.altitude = (
            self.fixed_wing_sub_pub.global_position_from_px4.altitude
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.latitude = (
            self.fixed_wing_sub_pub.global_position_from_px4.latitude
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.longitude = (
            self.fixed_wing_sub_pub.global_position_from_px4.longitude
        )

        # GPS速度是在ned下的
        self.fixed_wing_sub_pub.fw_states_from_mavros.global_vel_x = (
            self.fixed_wing_sub_pub.velocity_global_fused_from_px4.twist.linear.y
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.global_vel_y = (
            self.fixed_wing_sub_pub.velocity_global_fused_from_px4.twist.linear.x
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.global_vel_z = (
            -self.fixed_wing_sub_pub.velocity_global_fused_from_px4.twist.linear.z
        )

        # 以下为机体系和地面系之间的夹角、姿态角 IMU的坐标系为前左上坐标系
        # 添加符号转换到px4的坐标系,PX4所有值都在无人机体轴FRD（前右下）坐标系或者NED坐标系下表示
        self.fixed_wing_sub_pub.fw_states_from_mavros.roll_angle = (
            self.fixed_wing_sub_pub.att_angle_Euler[0]
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.pitch_angle = (
            - self.fixed_wing_sub_pub.att_angle_Euler[1]
        )
        if -self.fixed_wing_sub_pub.att_angle_Euler[2] + deg_2_rad(90.0) > 0:
            self.fixed_wing_sub_pub.fw_states_from_mavros.yaw_angle = (
                -self.fixed_wing_sub_pub.att_angle_Euler[2] + deg_2_rad(90.0)
            )  # 添加符号使增加方向相同，且大于PX4 90度, mavros坐标系中正东方向yaw为0，逆时针
        else:
            self.fixed_wing_sub_pub.fw_states_from_mavros.yaw_angle = (
                -self.fixed_wing_sub_pub.att_angle_Euler[2]
                + deg_2_rad(90.0)
                + deg_2_rad(360.0)
            )

        # 姿态四元数赋值
        self.att_angle[0] = self.fixed_wing_sub_pub.fw_states_from_mavros.roll_angle
        self.att_angle[1] = self.fixed_wing_sub_pub.fw_states_from_mavros.pitch_angle
        self.att_angle[2] = self.fixed_wing_sub_pub.fw_states_from_mavros.yaw_angle
        self.att_quat = euler_2_quat(self.att_angle)
        self.fixed_wing_sub_pub.fw_states_from_mavros.att_quater.w = self.att_quat[0]
        self.fixed_wing_sub_pub.fw_states_from_mavros.att_quater.x = self.att_quat[1]
        self.fixed_wing_sub_pub.fw_states_from_mavros.att_quater.y = self.att_quat[2]
        self.fixed_wing_sub_pub.fw_states_from_mavros.att_quater.z = self.att_quat[3]

        # 以下为ned坐标系下的位置，速度
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_pos_x = (
            self.fixed_wing_sub_pub.local_position_from_px4.pose.position.y
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_pos_y = (
            self.fixed_wing_sub_pub.local_position_from_px4.pose.position.x
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_pos_z = (
            -self.fixed_wing_sub_pub.local_position_from_px4.pose.position.z
        )

        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_vel_x = (
            self.fixed_wing_sub_pub.velocity_ned_fused_from_px4.twist.linear.y
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_vel_y = (
            self.fixed_wing_sub_pub.velocity_ned_fused_from_px4.twist.linear.x
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_vel_z = (
            -self.fixed_wing_sub_pub.velocity_ned_fused_from_px4.twist.linear.z
        )

        # 以下为body坐标系下的加速度，体轴系下的加速度与px4机体系的定义相同
        self.fixed_wing_sub_pub.fw_states_from_mavros.body_acc_x = (
            self.fixed_wing_sub_pub.imu.linear_acceleration.x
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.body_acc_y = (
            self.fixed_wing_sub_pub.imu.linear_acceleration.y
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.body_acc_z = (
            self.fixed_wing_sub_pub.imu.linear_acceleration.z
        )

        # 无人机的高度
        self.fixed_wing_sub_pub.fw_states_from_mavros.relative_alt = (
            self.fixed_wing_sub_pub.altitude_from_px4.relative
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ned_altitude = (
            self.fixed_wing_sub_pub.altitude_from_px4.local
        )

        # 无人机空速和地速
        self.fixed_wing_sub_pub.fw_states_from_mavros.air_speed = (
            self.fixed_wing_sub_pub.air_ground_speed_from_px4.airspeed
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.ground_speed = (
            self.fixed_wing_sub_pub.air_ground_speed_from_px4.groundspeed
        )

        # 风估计
        self.fixed_wing_sub_pub.fw_states_from_mavros.wind_estimate_x = (
            self.fixed_wing_sub_pub.wind_estimate_from_px4.twist.twist.linear.y
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.wind_estimate_y = (
            self.fixed_wing_sub_pub.wind_estimate_from_px4.twist.twist.linear.x
        )
        self.fixed_wing_sub_pub.fw_states_from_mavros.wind_estimate_z = (
            -self.fixed_wing_sub_pub.wind_estimate_from_px4.twist.twist.linear.z
        )

        # 电池状态
        # self.fixed_wing_sub_pub.fw_states_from_mavros.battery_current = (
        #     self.fixed_wing_sub_pub.battery_status_from_px4.current
        # )
        # self.fixed_wing_sub_pub.fw_states_from_mavros.battery_precentage = (
        #     self.fixed_wing_sub_pub.battery_status_from_px4.percentage
        # )
        # self.fixed_wing_sub_pub.fw_states_from_mavros.battery_voltage = (
        #     self.fixed_wing_sub_pub.battery_status_from_px4.voltage
        # )

        self.fixed_wing_states_pub.publish(
            self.fixed_wing_sub_pub.fw_states_from_mavros
        )

    def set_planeID(self, id:int):
        self.planeID = id
        if self.planeID == 0:
            self.uavID = "plane_0"
        elif self.planeID == 1:
            self.uavID = "plane_1"
        elif self.planeID == 2:
            self.uavID = "plane_2"
        elif self.planeID == 3:
            self.uavID = "plane_3"

    def msg_to_mavros(self):
        # 将期望值转换坐标系，并转化为四元数
        angle = np.zeros(3)
        quat = np.zeros(4)

        angle[0] = self.fixed_wing_sub_pub.cmd_from_controller.roll_angle_sp
        angle[1] = -self.fixed_wing_sub_pub.cmd_from_controller.pitch_angle_sp
        angle[2] = (
            -self.fixed_wing_sub_pub.cmd_from_controller.yaw_angle_sp + deg_2_rad(90.0)
        )

        quat = euler_2_quat(quat)

        self.fixed_wing_sub_pub.att_sp.type_mask = (
            7  # 1+2+4+64+128 body.rate_x,body.rate_y,body.rate_z thrust..
        )
        self.fixed_wing_sub_pub.att_sp.orientation.w = quat[0]
        self.fixed_wing_sub_pub.att_sp.orientation.x = quat[1]
        self.fixed_wing_sub_pub.att_sp.orientation.y = quat[2]
        self.fixed_wing_sub_pub.att_sp.orientation.z = quat[3]

        self.fixed_wing_sub_pub.att_sp.thrust = (
            self.fixed_wing_sub_pub.cmd_from_controller.throttle_sp
        )

    def srv_to_mavros(self):
        self.fixed_wing_sub_pub.mode_cmd._request_class.custom_mode = (
            self.fixed_wing_sub_pub.cmd_from_controller.cmd_mode
        )
        self.set_mode_service.call(self.fixed_wing_sub_pub.mode_cmd)

    def start(self):
        rospy.loginfo("pack_fw_states start")
        rate = rospy.Rate(50)
        self.ros_sub_pub()

        while not rospy.is_shutdown():

            # 持续向mavros发送，保证px4的链接
            self.pack_fw_states()
            rospy.loginfo("packing_plane{}_states".format(self.planeID))
            # self.msg_to_mavros()
            # self.srv_to_mavros()

            try:  # prevent garbage in console output when thread is killed
                rate.sleep()
            except rospy.ROSInterruptException as e:
                pass

        rospy.loginfo("pack_fw_states end")


if __name__ == "__main__":
    rospy.init_node("pack_fw_states", anonymous=True)
    pack_fw_states = PackFWStates()
    pack_fw_states.start()
